import torch
from transformers import GPT2Model, GPT2Config
import torch.optim as optim
import time

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义GPT-2模型的配置参数
config = GPT2Config(
    vocab_size=50257,  # GPT-2的词汇量
    n_embd=768,
    n_layer=12,
    n_head=12
)

# 根据定义的配置创建GPT-2模型
model = GPT2Model(config)
model = model.to(device)

# 设置优化器
optimizer = optim.Adam(model.parameters(), lr=1e-5)

# 指定批次大小
batch_size = 1 # 根据你的GPU容量调整这个值

# 初始化记录最大显存使用量
max_memory_used = 0

time_start = time.time()
for i in range(3):  # 简单示例：运行3个训练步骤
    # 创建随机输入ID
    input_ids = torch.randint(0, config.vocab_size, (batch_size, 512), dtype=torch.long).to(device)

    # 前向传播
    outputs = model(input_ids=input_ids)
    last_hidden_states = outputs.last_hidden_state

    # 为了简化，我们将所有激活的和作为损失
    loss = last_hidden_states.sum()

    # 反向传播
    loss.backward()

    # 更新模型
    optimizer.step()
    optimizer.zero_grad()

    # 检查并记录当前步骤后的最大显存使用
    current_memory = torch.cuda.max_memory_allocated(device)
    max_memory_used = max(max_memory_used, current_memory)
    
    # 重置显存峰值统计
    torch.cuda.reset_peak_memory_stats(device)
time_end = time.time()
# 打印记录的最大显存使用量
print(time_end-time_start)
print(f"Maximum GPU memory used: {max_memory_used / (1024**3):.15f} GB")
